#!/usr/bin/env python3
# -*- coding: utf-8 -*-

'ssh服务API'

import tkinter as tk
import os
import sys
import paramiko  # 导入 paramiko 库
import time      # 导入 time 库，用于延迟
import traceback
from tqdm import tqdm
import json
import logging.config
import shutil
from tkapp_bottom_o import AppBottom

class SSHService(object):

    def __init__(self):
        pass

    def setWriteFile(self):
        self.writeFile=SSHService.TqdmOutWriter(self._app_bottom.getOutText())
        pass

    def useBottom(self, bottom: AppBottom) -> AppBottom:
        self._app_bottom = bottom
        self.setWriteFile()
        pass

    def setup_logging(self, config_file='static/logging_config.json'):

        # # 判断是否是打包后的环境
        # if getattr(sys, 'frozen', False):
        #     # 打包后的路径
        #     base_path = sys._MEIPASS
        # else:
        #     # 开发环境的路径
        #     base_path = os.path.dirname(os.path.abspath(__file__))

        base_path = os.path.dirname(os.path.abspath(__file__))

        # 拼接配置文件路径
        config_path = os.path.join(base_path, config_file)

        # 检查文件是否存在
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"Logging config file not found: {config_path}")

        # 加载日志配置
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        logging.config.dictConfig(config)

        # # 检查文件是否存在
        # if not os.path.exists(config_path):
        #     raise FileNotFoundError(f"Logging config file not found: {config_path}")

        # with open(config_file, 'r') as f:
        #     config = json.load(f)
        # logging.config.dictConfig(config)
    
    # 使用日志记录器
    # setup_logging()
    # logger = logging.getLogger('ssh_service')
    # logger.debug('This is a debug message')
    # logger.info('This is an info message')
    # logger.error('This is an error message')

    # 通过用户名密码远程连接服务器
    def check_ssh_connection(self, host, port, username, password):
        """
        检查SSH连接状态的函数。
        :param host: 目标主机地址
        :param port: SSH端口，默认是22
        :param username: SSH用户名
        :param password: SSH密码
        :return: None
        """
        client = paramiko.SSHClient()  # 创建SSH客户端对象
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())  # 设置策略，自动添加主机密钥
        try:
            client.connect(hostname=host, port=port, username=username, password=password, timeout=5)  # 尝试连接
            print(f"成功连接到 {host}:{port}")
        except paramiko.AuthenticationException:
            print("认证失败，检查用户名或密码。")
        except paramiko.SSHException as e:
            print(f"SSH连接失败：{e}")
        except Exception as e:
            print(f"其他错误：{e}")
        finally:
            client.close()  # 关闭连接

    # 通过私钥字符串远程连接服务器
    def connectByPrivateKey(self, host, port, username, key): 
        # 也可以是存在于数据库中
        # key = """-----BEGIN RSA PRIVATE KEY-----
        # -----END RSA PRIVATE KEY-----"""
        # from io import StringIO
        # private_key = paramiko.RSAKey(file_obj=StringIO(key))

        # private_key = paramiko.RSAKey.from_private_key_file(r'/root/.ssh/id_rsa')
        private_key = paramiko.RSAKey.from_private_key_file(key) # 从本地文件获取
        
        # 创建SSH对象
        ssh = paramiko.SSHClient()
        # 允许连接不在know_hosts文件中的主机
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        # 连接服务器
        try:
            ssh.connect(hostname=host, port=port, username=username, pkey=private_key, timeout=5)
            # 执行命令
            # stdin, stdout, stderr = ssh.exec_command('df')
            # 获取命令结果
            # result = stdout.read()
            # print(result)
            print(f"成功连接到 {host}:{port}")
        except paramiko.AuthenticationException:
            print("认证失败，检查用户名或密码。")
        except paramiko.SSHException as e:
            print(f"SSH连接失败：{e}")
        except Exception as e:
            print(f"其他错误：{e}")
        finally:
            ssh.close()  # 关闭连接

    # 
    def getConnect(self, link):
        client = paramiko.SSHClient()  # 创建SSH客户端对象
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())  # 设置策略，自动添加主机密钥
        try:
            host = link[2]  # 替换为目标主机的IP地址
            port = 22             # 默认SSH端口
            username = link[6]  # 替换为你的用户名
            password = link[7]    # 替换为你的密码
            key_path = link[8]    # 私钥目录路径
            method = link[5]
            
            match method:
                case 1:
                    client.connect(hostname=host, port=port, username=username, password=password, timeout=5)  # 尝试连接
                    print(f"成功连接到 {host}:{port}")
                    return client
                case 2:
                    private_key = paramiko.RSAKey.from_private_key_file(key_path) # 从本地文件获取
                    # 允许连接不在know_hosts文件中的主机
                    # 连接服务器
                    client.connect(hostname=host, port=port, username=username, pkey=private_key, timeout=5)
                    # 执行命令
                    # stdin, stdout, stderr = client.exec_command('df')
                    # 获取命令结果
                    # result = stdout.read()
                    # print(result)
                    print(f"成功连接到 {host}:{port}")
                    return client
                case _: # _表示匹配到其他任何情况
                    print('无匹配模式')
        except paramiko.AuthenticationException:
            print("认证失败，检查用户名或密码。")
            client.close()  # 关闭连接
        except paramiko.SSHException as e:
            print(f"SSH连接失败：{e}")
            client.close()  # 关闭连接
        except Exception as e:
            print(f"其他错误：{e}")
            client.close()  # 
            # raise RuntimeError("连接异常")
        # finally:
        #     # client.close()  # 关闭连接
        #     return client

    # 通过 SSH 上传文件并显示进度
    def upload_file_with_progress(self, link, local_path, remote_path, progress_bar, ok=None):
        # # 创建 SSH 客户端
        # ssh = paramiko.SSHClient()
        # ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        # ssh.connect(hostname, port=port, username=username, password=password)

        ssh = self.getConnect(link)

        # 创建 SFTP 客户端
        sftp = ssh.open_sftp()

        # 远程目录是否存在
        remote_path_stat=True

        try:
            # 检查目录是否存在
            sftp.stat(remote_path)
            print(f"远程目录 '{remote_path}' 已存在。")
        except FileNotFoundError:
            # 目录不存在，创建目录
            try:
                sftp.mkdir(remote_path)
                print(f"远程目录 '{remote_path}' 创建成功。")
            except Exception as e:
                remote_path_stat=False
                print(f"创建远程目录 '{remote_path}' 时出错: {e}")
        except Exception as e:
                remote_path_stat=False
                print(f"检查远程目录 '{remote_path}' 时出错: {e}")

        if remote_path_stat==True and os.path.isfile(local_path):
            # 获取本地文件大小
            file_size = os.path.getsize(local_path)
            filename = os.path.basename(local_path)
            print(f"{local_path} 是一个文件，文件名是 {filename}")

            # # 使用 tqdm 显示进度条
            # with tqdm(total=file_size, unit='B', unit_scale=True, desc=os.path.basename(local_path)) as pbar:
            #     # 定义回调函数，用于更新进度条
            #     def callback(uploaded_bytes, total_bytes):
            #         pbar.update(uploaded_bytes - pbar.n)  # 更新进度条

            #     # 上传文件
            #     sftp.put(local_path, remote_path + "/" + filename, callback=callback)
            
            # 使用 tqdm 显示进度条
            # with TqdmToTkinter(total=file_size, unit='B', unit_scale=True, desc=os.path.basename(local_path), progress_bar=progress_bar) as pbar:
            #     def callback(uploaded_bytes, total_bytes):
            #         pbar.update(uploaded_bytes - pbar.n)  # 更新进度条

            #     # 上传文件
            #     sftp.put(local_path, remote_path + "/" + filename, callback=callback)
            #     print("任务完成")
            #     if ok != None:
            #         ok(ssh, filename, remote_path)
                    
            try:
                print("准备开始上传")
                with SSHService.TqdmToTkinter(writeFile=self.writeFile, total=file_size, unit='B', unit_scale=True, desc=os.path.basename(local_path), progress_bar=progress_bar) as pbar:
                    def callback(uploaded_bytes, total_bytes):
                        # print(uploaded_bytes)
                        pbar.update(uploaded_bytes - pbar.n)  # 更新进度条

                    # 上传文件
                    sftp.put(local_path, remote_path + "/" + filename, callback=callback)

                    time.sleep(1)  # 延时1秒
                    print("\n任务完成\n")
                    pbar.reset_progressbar()
                    if ok != None:
                        ok(ssh, filename, remote_path)
            except Exception as e:
                print(f"其他错误：{e}")
                print("详细的堆栈信息:")
                # traceback.print_exc()
                error_message = traceback.format_exc()
                print(error_message)

                
        elif os.path.isdir(local_path):
            print(f"{local_path} 是一个目录，不是文件")
            self.compressAndUpoad(link, local_path, remote_path, progress_bar)
        else:
            print(f"{local_path} 既不是文件也不是目录")

        # 关闭连接
        sftp.close()
        ssh.close()

    # 示例使用
    # if __name__ == "__main__":
    #     hostname = 'your.server.com'  # 替换为服务器地址
    #     port = 22  # SSH 端口，默认是 22
    #     username = 'your_username'  # 替换为用户名
    #     password = 'your_password'  # 替换为密码
    #     local_path = 'path/to/your/local/file.txt'  # 替换为本地文件路径
    #     remote_path = 'path/to/your/remote/file.txt'  # 替换为远程文件路径

    #     upload_file_with_progress(hostname, port, username, password, local_path, remote_path)

    # 压缩目录后上传
    def compressAndUpoad(self, link, local_path, remote_path, progress_bar):
        dirname = os.path.basename(local_path)
        print(f"{local_path} 是一个目录，目录名是 {dirname}")
        # 压缩
        compressFilePath=self.compress_directory(local_path, dirname)
        # 压缩完成，上传
        self.upload_file_with_progress(link, compressFilePath, remote_path, progress_bar, self.unzip)

    # 上传完成后解压缩
    def unzip(self, client, filename, remote_path):

        file_name, file_extension = os.path.splitext(filename)

        # 要执行的命令列表
        commands = [
            f'cd {remote_path}',                        # 切换到目录
            # 'ls -l',                                  # 列出目录内容
            f'sudo rm -rf ./{file_name}',
            f'sudo unzip {filename} -d ./{file_name}',  # 解压
            f'pwd',
            'ls -l'                                     # 再次列出目录内容
        ]

        # 将命令用 && 连接
        command = ' && '.join(commands)

        self.send_command_client(command, client)


    # 上传整个目录，遍历目录中的所有文件，并逐个上传：
    # def upload_directory_with_progress(hostname, port, username, password, local_dir, remote_dir):
    def upload_directory_with_progress(self, link, local_dir, remote_dir, progress_bar):
        # 创建 SSH 客户端
        # ssh = paramiko.SSHClient()
        # ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        # ssh.connect(hostname, port=port, username=username, password=password)

        ssh = self.getConnect(link)

        # 创建 SFTP 客户端
        sftp = ssh.open_sftp()

        # 遍历本地目录
        for root, dirs, files in os.walk(local_dir):
            # 计算远程目录路径
            relative_path = os.path.relpath(root, local_dir)
            remote_path = os.path.join(remote_dir, relative_path)

            # 在远程服务器上创建目录
            try:
                sftp.mkdir(remote_path)
            except IOError:
                pass  # 如果目录已存在，则忽略错误

            # 上传文件
            for file in files:
                local_file_path = os.path.join(root, file)
                remote_file_path = os.path.join(remote_path, file)

                # 获取本地文件大小
                file_size = os.path.getsize(local_file_path)

                # # 使用 tqdm 显示进度条
                # with tqdm(total=file_size, unit='B', unit_scale=True, desc=file) as pbar:
                #     def callback(uploaded_bytes, total_bytes):
                #         pbar.update(uploaded_bytes - pbar.n)  # 更新进度条

                #     # 上传文件
                #     sftp.put(local_file_path, remote_file_path, callback=callback)

                # 使用 tqdm 显示进度条
                with SSHService.TqdmToTkinter(total=file_size, unit='B', unit_scale=True, desc=file, progress_bar=progress_bar) as pbar:
                    def callback(uploaded_bytes, total_bytes):
                        pbar.update(uploaded_bytes - pbar.n)  # 更新进度条

                    # 上传文件
                    sftp.put(local_file_path, remote_file_path, callback=callback)
                    

        # 关闭连接
        sftp.close()
        ssh.close()

    # 示例使用
    # if __name__ == "__main__":
    #     hostname = 'your.server.com'  # 替换为服务器地址
    #     port = 22  # SSH 端口，默认是 22
    #     username = 'your_username'  # 替换为用户名
    #     password = 'your_password'  # 替换为密码
    #     local_dir = 'path/to/your/local/directory'  # 替换为本地目录路径
    #     remote_dir = 'path/to/your/remote/directory'  # 替换为远程目录路径

    #     upload_directory_with_progress(hostname, port, username, password, local_dir, remote_dir)

    # 压缩目录
    def compress_directory(self, source_dir, dirname):
        # output_path = os.path.join(source_dir, dirname)
        shutil.make_archive(source_dir, 'zip', root_dir=source_dir)
        compressFilePath=f"{source_dir}.zip"
        print(f'目录已压缩为: {compressFilePath}')
        return compressFilePath

    # 示例
    # compress_directory("/path/to/directory", "archive_name")

    # 发送命令
    def send_command(self, command, link):
        try:
            # 连接到远程服务器
            client = self.getConnect(link)

            print("执行命令")
            print(command)
            # 执行命令
            stdin, stdout, stderr = client.exec_command(command)

            # 读取命令输出
            output = stdout.read().decode()
            error = stderr.read().decode()

            # 打印输出
            print("Output:")
            print(output)

            if error:
                print("Error:")
                print(error)

        finally:
            # 关闭连接
            client.close()
            print("Connection closed.")

    # 发送命令
    def send_command_client(self, command, client):
        try:
            print("执行命令")
            print(command)
            # 执行命令
            stdin, stdout, stderr = client.exec_command(command)

            # 读取命令输出
            output = stdout.read().decode()
            error = stderr.read().decode()

            # 打印输出
            print("Output:")
            print(output)

            if error:
                print("Error:")
                print(error)

        except Exception as e:
            print(f"发送命令错误：{e}") 

    '''
    TqdmToTkinter 类：
    继承自 tqdm，重写 update 方法。
    在每次更新时，将进度信息传递给 tkinter 的 ttk.Progressbar
    '''
    class TqdmToTkinter(tqdm):
        def __init__(self, *args, **kwargs):
            self.progress_bar = kwargs.pop('progress_bar', None)  # 获取 tkinter 进度条
            self.writeFile = kwargs.pop('writeFile', None)  # 获取 tkinter 进度条
            super().__init__(file=self.writeFile, *args, **kwargs)

        def update(self, n=1):
            result = super().update(n)  # 调用父类的 update 方法
            if self.progress_bar:
                # 更新 tkinter 进度条
                self.progress_bar['value'] = self.n  # 当前进度
                self.progress_bar['maximum'] = self.total  # 总进度
                self.progress_bar.update_idletasks()  # 刷新界面
            return result
        
        def reset_progressbar(self):
            # 将Progressbar的进度设置为0
            self.progress_bar.config(value=0)
        
    # 自定义一个类，用于重定输出
    class TqdmOutWriter:
        def __init__(self, text_widget):
            self.text_widget = text_widget

        def write(self, message):
            if len(message) > 0:
                self.delete_last_line()
                # self.text_widget.mark_set("insert", "insert linestart")
                self.text_widget.insert(tk.END, message)  # 将输出插入到 Text 组件中
                self.text_widget.see(tk.END)  # 自动滚动到末尾

        def flush(self):
            pass  # 需要实现 flush 方法，但可以为空

        def delete_last_line(self):
            # 获取最后一行的起始和结束索引
            last_line_start = self.text_widget.index("end-2l linestart")
            last_line_end = self.text_widget.index("end-1c")
            print(f'{last_line_start}-{last_line_end}')
            # 删除最后一行
            self.text_widget.delete(last_line_start, last_line_end)

    def long_running_task(self, progress_bar):
        # 模拟一个耗时的任务
        for _ in SSHService.TqdmToTkinter(range(100), progress_bar=progress_bar):
            time.sleep(0.1)  # 模拟任务延迟





